from __future__ import absolute_import

import os
import os.path as osp
import sys

from .iotools import mkdir_if_missing


class Logger(object):
    """
    Write console output to external text file.
    Code imported from https://github.com/Cysu/open-reid/blob/master/reid/utils/logging.py.
    """

    def __init__(self, fpath=None):
        self.console = sys.stdout
        self.file = None
        if fpath is not None:
            file_path = osp.dirname(fpath)
            mkdir_if_missing(file_path)
            self.file = open(fpath, 'w')

    def __del__(self):
        self.close()

    def __enter__(self):
        pass

    def __exit__(self, *args):
        self.close()

    def write(self, msg):
        self.console.write(msg)
        if self.file is not None:
            self.file.write(msg)

    def flush(self):
        self.console.flush()
        if self.file is not None:
            self.file.flush()
            os.fsync(self.file.fileno())

    def close(self):
        self.console.close()
        if self.file is not None:
            self.file.close()


class RankLogger(object):
    """
    RankLogger records the rank1 matching accuracy obtained for each
    test dataset at specified evaluation steps and provides a function
    to show the summarized results, which are convenient for analysis.
    Args:
    - source_names (list): list of strings (names) of source datasets.
    - target_names (list): list of strings (names) of target datasets.
    """

    def __init__(self, dataset_name):
        self.dataset_name = dataset_name
        self.logger = {dataset_name: {'epoch': [], 'rank1': []}}

    def write(self, name, epoch, rank1):
        self.logger[name]['epoch'].append(epoch)
        self.logger[name]['rank1'].append(rank1)

    def show_summary(self):
        print(f'=> Show {self.dataset_name} performance summary')
        for epoch, rank1 in zip(self.logger[self.dataset_name]['epoch'], self.logger[self.dataset_name]['rank1']):
            print('- epoch {}\t rank1 {:.1%}'.format(epoch, rank1))


class AccLogger(object):
    """
    RankLogger records the rank1 matching accuracy obtained for each
    test dataset at specified evaluation steps and provides a function
    to show the summarized results, which are convenient for analysis.
    Args:
    - source_names (list): list of strings (names) of source datasets.
    - target_names (list): list of strings (names) of target datasets.
    """

    def __init__(self, dataset_name):
        self.dataset_name = dataset_name
        self.logger = {dataset_name: {'epoch': [], 'acc': []}}

    def write(self, name, epoch, acc):
        self.logger[name]['epoch'].append(epoch)
        self.logger[name]['acc'].append(acc)

    def show_summary(self):
        print('=> Show performance summary')
        for epoch, acc in zip(self.logger[self.dataset_name]['epoch'], self.logger[self.dataset_name]['acc']):
            print('- epoch {}\t acc {:.1%}'.format(epoch, acc))
